- Update test_datasets.py to use ForecasterModule instead of GraphLAM
- Update test_plotting.py to use ForecasterModule instead of GraphLAM
- Fix interior_mask_bool property shape (1,) -> (N,) for correct loss masking
- Fix all_gather_cat to handle single-device runs without incorrect dim collapse

…r hierarchy

- Replace opaque argparse.Namespace with explicit keyword arguments in
  StepPredictor, BaseGraphModel, BaseHiGraphModel, GraphLAM, HiLAM,
  and HiLAMParallel __init__ methods
- Reorder methods in step_predictor.py: forward/expand_to_batch now
  appear before clamping methods
- Update all instantiation sites (train_model.py, test_training.py,
  test_prediction_model_classes.py) to pass explicit kwargs
- HiLAM helper methods (make_same/up/down_gnns) now use self.hidden_dim
  and self.hidden_layers instead of args parameter

Addresses review comments on PR mllam#208.

- Rename border to boundary in Forecaster
- Pass Forecaster object to ForecasterModule init instead of Predictor
- Remove inline imports in ForecasterModule
- Move loss-related pred_std logic fully into ForecasterModule
- Delete obsolete test_refactored_hierarchy.py

…phModel.__init__

Co-authored-by: Joel Oskarsson <joel.oskarsson@outlook.com>

- Add predicts_std property to StepPredictor, Forecaster and ARForecaster
  so ForecasterModule can query the forecaster instead of taking output_std
  as a separate constructor argument
- Remove output_std parameter from ForecasterModule; use
  self._forecaster.predicts_std throughout
- Move fallback per_var_std logic out of forecast_for_batch into each
  step method so pred_std is None before fallback, enabling direct None
  checks instead of hparam checks
- Replace len(datastore.boundary_mask) with datastore.num_grid_points in
  StepPredictor to avoid relying on boundary_mask
- Move get_state_feature_weighting and ARForecaster inline imports to
  module-level imports in forecaster_module.py and train_model.py
- Fix statement ordering in StepPredictor.__init__ so register_buffer for
  grid_static_features appears directly after building the tensor
- Replace dict+loop pattern for registering state_mean/state_std buffers
  with two direct register_buffer calls
- Remove all internal Item N checklist references from comments
- Remove TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD env var hack; pass
  weights_only=False explicitly to load_from_checkpoint calls and
  weights_only=True to torch.load in test_graph_creation.py
- Add test_step_predictor_no_static_features to verify models initialise
  and run correctly when the datastore returns None for static features
- Fix graph= -> graph_name= and model.forecaster -> model._forecaster in
  tests to match current API

…r_batch

Makes the forecasting path tolerant to batch-folded execution so that
future ensemble generation can fold (S, B) into (S*B) before calling
ARForecaster, without any changes to ARForecaster or StepPredictor.

Prediction is kept folded through the existing deterministic logging and
aggregation paths so all dim assumptions in training_step, validation_step,
and test_step remain correct. Unfolding to (*leading, T, N, F) is deferred
to ensemble-specific subclasses (e.g. EnsForecasterModule).

Adds test_fold_unfold_equivalence to confirm ARForecaster's rollout is
rank-transparent under a pre-entry fold.

…or_batch

…stic models

- Port PropagationNet as InteractionNet subclass (mean aggr, sender residual
  in messages, aggregation residual in forward)
- Add --vertical_propnets CLI flag to select PropagationNet for grid-mesh
  and vertical message passing edges
- Wire flag through model hierarchy: BaseGraphModel (g2m/m2g),
  BaseHiGraphModel (mesh init), HiLAM (up GNNs)
- Add 13 tests covering unit behavior and backward compatibility